{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "from hmmlearn.hmm import MultinomialHMM\n",
    "from sklearn.metrics import classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "term_dict = {}\n",
    "word_dict = {' ':0}\n",
    "hidden_dict = {'B':0,'M':1,'E':2,'S':3}\n",
    "total_obser_seq = []\n",
    "total_hidden_seq = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_data_set(data, term_dict, word_dict, hidden_dict):\n",
    "    total_obser_seq = []\n",
    "    total_hidden_seq = []\n",
    "    for _, i in enumerate(data):\n",
    "        line = i.strip(' \\n').split('  ')\n",
    "        obser_seq = list(''.join(line))\n",
    "        hidden_seq = ''\n",
    "        for term in line:\n",
    "            # BMES\n",
    "            if len(term)==1:\n",
    "                hidden_seq+='S'\n",
    "            elif len(term)==0:\n",
    "                continue\n",
    "            else:\n",
    "                hid = 'B'+'M'*(len(term)-2)+'E'\n",
    "                hidden_seq+=hid    \n",
    "            # 词表\n",
    "            term_dict[term] = len(term_dict)\n",
    "        for j in range(len(obser_seq)):\n",
    "            if word_dict.get(obser_seq[j],'NO') != 'NO':\n",
    "                obser_seq[j]=word_dict[obser_seq[j]]\n",
    "            else:\n",
    "                length = len(word_dict)\n",
    "                # print(word_dict, obser_seq[j], length)\n",
    "                word_dict[obser_seq[j]] = length\n",
    "                obser_seq[j] = length\n",
    "        hidden_seq = [hidden_dict[i] for i in hidden_seq]\n",
    "        if obser_seq==[]:\n",
    "            continue           \n",
    "        else:\n",
    "            total_obser_seq.append(obser_seq)\n",
    "            total_hidden_seq.append(hidden_seq)\n",
    "    return total_hidden_seq, total_obser_seq, term_dict, word_dict, hidden_dict\n",
    "\n",
    "    \n",
    "    \n",
    "# f = open(r'../data/segment sentence data/msr_training.utf8', 'r', encoding='utf-8')\n",
    "# data = f.readlines()\n",
    "# h, o, term_dict, word_dict, hidden_dict = create_data_set(data, term_dict, word_dict, hidden_dict)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = [r'../data/segment sentence data/msr_training.utf8', r'../data/segment sentence data/pku_training.utf8']\n",
    "for file in files:\n",
    "    f = open(file, 'r', encoding='utf-8')\n",
    "    data = f.readlines()\n",
    "    h, o, term_dict, word_dict, hidden_dict = create_data_set(data, term_dict, word_dict, hidden_dict)\n",
    "    total_obser_seq.extend(o)\n",
    "    total_hidden_seq.extend(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, x in enumerate(total_obser_seq):\n",
    "    if x==[]:\n",
    "        print('%s!'%i)\n",
    "        print(total_hidden_seq[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_init_array(total_hidden_seq, hidden_dict):\n",
    "    pi = [0 for i in hidden_dict]\n",
    "    for _ in total_hidden_seq:\n",
    "        pi[_[0]]+=1\n",
    "    stat_array = pi\n",
    "    prob_array = [i/sum(pi) for i in pi]\n",
    "    return prob_array\n",
    "\n",
    "\n",
    "def cal_transfer_matrix(total_hidden_seq, hidden_dict):\n",
    "    matrix = np.zeros([len(hidden_dict),len(hidden_dict)])\n",
    "    for i in total_hidden_seq:\n",
    "        for j in range(len(i)-1):\n",
    "            matrix[i[j],i[j+1]]+=1\n",
    "    for i in range(len(matrix)):\n",
    "        matrix[i] = np.divide(matrix[i], sum(matrix[i]))\n",
    "    return matrix\n",
    "\n",
    "def cal_emmit_matrix(total_obser_seq, total_hidden_seq, hidden_dict, word_dict):\n",
    "    matrix = np.zeros([len(hidden_dict), len(word_dict)])\n",
    "    for hid_seq, obs_seq in zip(total_hidden_seq, total_obser_seq):\n",
    "        for hid, obs in zip(hid_seq, obs_seq):\n",
    "            matrix[hid][obs]+=1\n",
    "    for i in range(len(matrix)):\n",
    "        matrix[i] = np.divide(matrix[i], sum(matrix[i]))\n",
    "    return matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "a = [i for i in range(len(total_obser_seq))]\n",
    "random.shuffle(a)\n",
    "trainlist = a[:round(len(total_obser_seq)*0.7)]\n",
    "testlist = a[round(len(total_obser_seq)*0.7):]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainobs = [total_obser_seq[i] for i in trainlist]\n",
    "trainhid = [total_hidden_seq[i] for i in trainlist]\n",
    "testobs = [total_obser_seq[i] for i in testlist]\n",
    "testhid = [total_hidden_seq[i] for i in testlist]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "seg_hmm = MultinomialHMM(n_components=4)\n",
    "seg_hmm.startprob_ = cal_init_array(trainhid, hidden_dict)\n",
    "seg_hmm.transmat_ = cal_transfer_matrix(trainhid, hidden_dict)\n",
    "seg_hmm.emissionprob_ = cal_emmit_matrix(trainobs, trainhid, hidden_dict, word_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hmm_decode(hmm_model, word_list, word_dict, hidden_dict):\n",
    "    \n",
    "    input = np.reshape(np.array(word_list),(len(np.array(word_list)),1))\n",
    "    result = hmm_model.decode(input, algorithm='viterbi')   \n",
    "    sentence = [word_dict[i] for i in word_list]\n",
    "    segment = [hidden_dict[i] for i in result[1]]\n",
    "    return sentence, segment\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "id2word = {v:k for k,v in word_dict.items()}\n",
    "id2hidden = {v:k for k,v in hidden_dict.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_predict = []\n",
    "total_truth = []\n",
    "for i in range(len(testhid)):\n",
    "    try:\n",
    "        if len(testobs[i])!=len(testhid[i]): \n",
    "            print(testhid[i],testobs[i]) \n",
    "            raise ValueError('wrong num')\n",
    "        else:\n",
    "            sent, seg = hmm_decode(seg_hmm, testobs[i], id2word, id2hidden)\n",
    "            total_predict.extend(seg)\n",
    "            total_truth.extend([id2hidden[j] for j in testhid[i]])\n",
    "    except:\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1762095"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum([len(i) for i in testhid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           B       0.81      0.87      0.83    552120\n",
      "           E       0.80      0.87      0.83    552120\n",
      "           M       0.64      0.44      0.52    167272\n",
      "           S       0.84      0.78      0.81    490583\n",
      "\n",
      "    accuracy                           0.80   1762095\n",
      "   macro avg       0.77      0.74      0.75   1762095\n",
      "weighted avg       0.80      0.80      0.80   1762095\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(total_truth, total_predict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = '十几年来，已有美国、日本、法国、英国、德国、芬兰、意大利、新加坡、香港、台湾等十几个国家和地区的境外企业进入中国进行工程总承包或工程分包。'\n",
    "a,b = hmm_decode(seg_hmm, [word_dict[i] for i in input], id2word, id2hidden)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['十几年来', '，', '已', '有', '美国', '、', '日本', '、', '法国', '、', '英国', '、', '德国', '、', '芬兰', '、', '意大利', '、', '新', '加坡', '、', '香港', '、', '台湾', '等', '十几个', '国家', '和', '地区', '的', '境外', '企业', '进入', '中国', '进行', '工程', '总', '承包', '或', '工程', '分包', '。']\n"
     ]
    }
   ],
   "source": [
    "def segment_result(sentence, seg_tag):\n",
    "    if len(sentence) != len(seg_tag):\n",
    "        return 'diff length between inputs'\n",
    "    \n",
    "    result = []\n",
    "    temp_term = ''\n",
    "    for index, tag in enumerate(seg_tag):\n",
    "        if tag=='E' or tag == 'S':\n",
    "            temp_term += sentence[index]\n",
    "            result.append(temp_term)\n",
    "            temp_term = ''\n",
    "        else:\n",
    "            temp_term += sentence[index]\n",
    "    \n",
    "    return result\n",
    "\n",
    "print(segment_result(a,b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "np.save('hmm_startprob.npy', seg_hmm.startprob_)\n",
    "np.save('hmm_transmat.npy', seg_hmm.transmat_)\n",
    "np.save('hmm_emissionmat.npy', seg_hmm.emissionprob_)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6d46af94c2bbce495f1e668725902fa517c90b1782bcfe2fce0dd9868df553d3"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
